#ifndef MLP_H
#define MLP_H
#include <Tim_STL.h>
#include <Eigen/Dense>

#define delta 0.0000001
#define ACT ReLU
#define LOSS MSE_Loss

class Layer
{
	public:
		Eigen::MatrixXd W;
		Eigen::VectorXd B;
		Eigen::VectorXd a;
		Eigen::VectorXd z;
		Eigen::VectorXd X;
		int in_size, out_size;
		double (*activate)(double x);
		double (*gradient)(double x);

		Layer(int input_size = 0, int output_size = 0,
			  double (*A)(double x) = nullptr,
			  double (*G)(double x) = nullptr
			  );
		Eigen::VectorXd& operator()(Eigen::VectorXd const& x);
		friend std::ostream& operator<<(std::ostream& out, Layer& layer);
};

class SGD
{
	public:
		double LR;
		double regular_rate;
		Eigen::VectorXd error;
		Eigen::MatrixXd last_W;
		Eigen::VectorXd Y;
		bool should;

		double (*loss)(Eigen::VectorXd const& y, Eigen::VectorXd const& y_hat);
		SGD(
			double lr, double lambda,
			double (*lossfuc)(Eigen::VectorXd const& y, Eigen::VectorXd const& y_hat)
			):
			LR(lr), error(1),
			regular_rate(lambda),
			loss(lossfuc), should(true) { error[0] = 0.0; };
		void operator()(Layer& layer);
		void set(Layer& layer);
		void ground_truth(Eigen::VectorXd const& y) { Y = y; };
		void step_zero();
};

class MLP: public tim::list<Layer>
{
	public:
		MLP(int shape[], int n);

		void add_BatchNormalization(int i);
		void add_Softmax();
		void add_Dropout(double rate);

		Eigen::VectorXd forward(Eigen::VectorXd const& x);

		template<typename VST>
		double step(
					VST& optmizer,
					double (*lossfuc)(Eigen::VectorXd const& y, Eigen::VectorXd const& y_hat),
					Eigen::VectorXd const& x,
					Eigen::VectorXd const& y
				);

		template<typename VST>
		void train(
				int epoch, VST& optmizer, int show_interval, int save_interval,
				double (*lossfuc)(Eigen::VectorXd const& y, Eigen::VectorXd const& y_hat),
				Eigen::VectorXd features[], Eigen::VectorXd labels[], int size_,
				double (*accfuc)(Eigen::VectorXd const& y, Eigen::VectorXd const& y_hat),
				void (*save)(MLP const& mlp)
				);

		friend std::ostream& operator<<(std::ostream& out, MLP& mlp);
};


#endif
